import gymnasium as gym
import torch
import numpy as np

agent = torch.load('model.pth')

if __name__ == '__main__':
    env = gym.make('CartPole-v1', render_mode="human")
    observation, info = env.reset()
    episode_over = False
    while True:
        action_probe = agent(torch.from_numpy(observation).float())
        action = np.random.choice(len(action_probe), p=action_probe.detach().numpy())
        observation, reward, terminated, truncated, info = env.step(action)
        episode_over = terminated or truncated
